#include "key_point.h"
#include "common\common.hpp"
#include "common\gpu_utility.h"
#include <assert.h>

namespace key_point
{
	std::vector<char> get_file_data(std::string file_path)
	{
		std::ifstream file(file_path, std::ios_base::binary);
		assert(file.good());
		file.seekg(0, std::ios_base::end);
		int file_len = file.tellg();
		file.seekg(0, std::ios_base::beg);

		std::vector<char> data(file_len, 0);
		file.read(data.data(), file_len);

		file.close();
		return data;
	}

	void decrypt_data(std::vector<char> data)
	{
		char trans = 55;
		for (int i = 0; i < data.size(); ++i)
		{
			data[i] -= trans;
			trans += 11;
		}
	}

	bool decryptFile(std::string file_path, st_encrypt_info& info, std::vector<char>& engine_data)
	{
		std::vector<char> file_data = get_file_data(file_path);
		decrypt_data(file_data);
		if (file_data.size() <= sizeof(st_encrypt_info))
		{
			std::cout << "file data too small!\n";
			return false;
		}
		info = *((st_encrypt_info*)file_data.data());
		if (info.engine_size + sizeof(st_encrypt_info) != file_data.size())
		{
			std::cout << "file data is bad format!\n";
			return false;
		}

		std::string gpu_model = GetGPUModel();
		std::string build_gpu_model = info.gpu_model;
		if (gpu_model != build_gpu_model)
		{
			std::cout << "gpu model warning: build is: " << build_gpu_model << " target is: " << gpu_model << std::endl;
		}

		int index = sizeof(st_encrypt_info);
		engine_data.resize(info.engine_size);
		memcpy(engine_data.data(), file_data.data() + index, info.engine_size);
		return true;
	}


	bool set_gpu_id(const int id = 0)
	{
		cudaError_t status = cudaSetDevice(id);
		if (status != cudaSuccess)
		{
			std::cout << "gpu id :" + std::to_string(id) + " not exist !" << std::endl;
			return false;
		}
		return true;
	}

	CKeyPoint::CKeyPoint()
	{
		cudaStreamCreate(&m_stream);
	}

	CKeyPoint::~CKeyPoint()
	{
		cudaStreamDestroy(m_stream);
		Reset();
	}





	void CKeyPoint::Reset()
	{
		if (m_context)
		{
			m_context->destroy();
		}
		if (m_engine)
		{
			m_engine->destroy();
		}
		for (int i = 0; i < m_v_dev_buffer.size(); ++i)
		{
			if (m_v_dev_buffer[i])
			{
				cudaFree(m_v_dev_buffer[i]);
				m_v_dev_buffer[i] = 0;
			}
		}
		m_v_dev_buffer.clear();
	}

	bool CKeyPoint::Init(int gpu_id, const char* gdats_path)
	{
		std::string file_path = gdats_path;
		std::vector<char> v_engine_data;
		if (!set_gpu_id(gpu_id) || !decryptFile(file_path, m_info, v_engine_data))
		{
			return false;
		}
		return LoadEngine(v_engine_data);
	}
	bool CKeyPoint::LoadEngine(std::vector<char> v_engine_data)
	{
		Reset();

		nvinfer1::IRuntime* trtRuntime;
		trtRuntime = nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger());
		m_engine = trtRuntime->deserializeCudaEngine(v_engine_data.data(), v_engine_data.size(), nullptr);
		if (m_engine == nullptr)
		{
			std::cout << "deserialize fail" << std::endl;
			return false;
		}
		m_context = m_engine->createExecutionContext();
		assert(m_context != nullptr);

		assert(m_engine->getNbBindings() == 2);
		int nbBindings = m_engine->getNbBindings();
		m_v_dev_buffer.resize(nbBindings);
		m_v_dev_buf_size.resize(nbBindings);

		for (int i = 0; i < nbBindings; ++i) {
			nvinfer1::Dims dims = m_engine->getBindingDimensions(i);
			nvinfer1::DataType dtype = m_engine->getBindingDataType(i);
			int64_t totalSize = volume(dims) * 1 * getElementSize(dtype);
			m_v_dev_buf_size[i] = totalSize;
			//std::cout << "binding" << i << ": " << totalSize << std::endl;
			cudaMalloc(&m_v_dev_buffer[i], totalSize);
		}
		m_out_size[0] = m_v_dev_buf_size[1] / sizeof(float) / m_info.batch_size;
		
		m_num_anchors = 0;
		for (int i = 0;i<m_info.stride_num;i++)
		{
			m_num_anchors += (m_info.image_height / m_info.strides[i]) * (m_info.image_width / m_info.strides[i]);
		}
		return true;
	}
	
	cv::Mat Data2Mat(image_t & data)
	{
		using namespace cv;
		int rows = data.h, cols = data.w;
		Mat mat = Mat::zeros(rows, cols, CV_32FC3);
		if (data.rgbf[0] && data.rgbf[1] && data.rgbf[2])
		{
			float * p_r = data.rgbf[0];
			float * p_g = data.rgbf[1];
			float * p_b = data.rgbf[2];
			for (int i = 0; i < rows; ++i)
			{
				for (int j = 0; j < cols; ++j)
				{
					Vec3f val;
					val[2] = p_r[i*cols + j] * 255.f;
					val[1] = p_g[i*cols + j] * 255.f;
					val[0] = p_b[i*cols + j] * 255.f;
					mat.at<Vec3f>(i, j) = val;
				}
			}
		}
		else if (data.rgb16[0] && data.rgb16[1] && data.rgb16[2])
		{
			unsigned short * p_r = data.rgb16[0];
			unsigned short * p_g = data.rgb16[1];
			unsigned short * p_b = data.rgb16[2];
			for (int i = 0; i < rows; ++i)
			{
				for (int j = 0; j < cols; ++j)
				{
					Vec3f val;
					val[2] = float(p_r[i*cols + j] >> 8);
					val[1] = float(p_g[i*cols + j] >> 8);
					val[0] = float(p_b[i*cols + j] >> 8);
					mat.at<Vec3f>(i, j) = val;
				}
			}
		}
		else if (data.rgb8[0] && data.rgb8[1] && data.rgb8[2])
		{
			unsigned char * p_r = data.rgb8[0];
			unsigned char * p_g = data.rgb8[1];
			unsigned char * p_b = data.rgb8[2];
			for (int i = 0; i < rows; ++i)
			{
				for (int j = 0; j < cols; ++j)
				{
					Vec3f val;
					val[2] = float(p_r[i*cols + j]);
					val[1] = float(p_g[i*cols + j]);
					val[0] = float(p_b[i*cols + j]);
					mat.at<Vec3f>(i, j) = val;
				}
			}
		}
		else
		{
			assert(0);
		}
		return mat;
	}



	std::vector<float> CKeyPoint::PrepareImage(std::vector<cv::Mat> vec_img)
	{
		int batch = m_info.batch_size;
		int model_width = m_info.image_width;
		int model_height = m_info.image_height;
		int model_channel = m_info.input_channel;

		std::vector<float> result(batch * model_width * model_height * model_channel);
		float* data = result.data();
		int index = 0;

		for (const cv::Mat& src_img : vec_img)
		{
			if (!src_img.data)
				continue;

			float ratio = __min(float(model_width) / float(src_img.cols), float(model_height) / float(src_img.rows));

			cv::Mat rsz_img;
			cv::Mat out_img;
			cv::resize(src_img, rsz_img, cv::Size(), ratio, ratio);

			float padd_w = (model_width - src_img.cols * ratio) / 2;
			float padd_h = (model_height - src_img.rows * ratio) / 2;

			float top = padd_h;
			float bottom = padd_h;
			float left = padd_w;
			float right = padd_w;

			cv::copyMakeBorder(rsz_img, rsz_img, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
			//cv::imshow("", rsz_img);
			//cv::waitKey(0);
			cv::resize(rsz_img, rsz_img, cv::Size(model_width, model_height));

			rsz_img.convertTo(rsz_img, CV_32FC3, 1.0 / 255, 0);

			int channelLength = model_width * model_height;
			std::vector<cv::Mat> split_img = {
				cv::Mat(model_width, model_height, CV_32FC1, data + channelLength * (index + 2)),
				cv::Mat(model_width, model_height, CV_32FC1, data + channelLength * (index + 1)),
				cv::Mat(model_width, model_height, CV_32FC1, data + channelLength * index)
			};
			index += 3;
			cv::split(rsz_img, split_img);
		}
		return result;
	}

	void CKeyPoint::Detect(st_dlm_data* p_imgs, int num, std::vector<std::vector<DetectRes>> &vv_results)
	{
		auto t_start = std::chrono::high_resolution_clock::now();

		//Ԥ���� ͼ������================================================================
		int net_w = m_info.image_width,
			net_h = m_info.image_height,
			net_c = m_info.input_channel;

		int batchSize = num;
		std::vector<cv::Mat> vec_Mat;
		for (int i = 0; i < num; ++i)
		{
			st_dlm_data& data = p_imgs[i];
			image_t & img = data.img;
			vec_Mat.push_back(Data2Mat(img));
		}

		std::vector<float> curInput = PrepareImage(vec_Mat);

		cudaMemcpyAsync(m_v_dev_buffer[0], curInput.data(), m_v_dev_buf_size[0], cudaMemcpyHostToDevice, m_stream);
		m_context->enqueueV2(m_v_dev_buffer.data(), m_stream, nullptr);

		for (int i = 0; i < 1; ++i)
		{
			if (m_v_out_buf[i].size() != m_out_size[i] * m_info.batch_size)
			{
				m_v_out_buf[i].resize(m_out_size[i] * m_info.batch_size);
			}
			cudaMemcpyAsync(m_v_out_buf[i].data(), m_v_dev_buffer[1 + i], m_v_dev_buf_size[1 + i], cudaMemcpyDeviceToHost, m_stream);
		}
		cudaStreamSynchronize(m_stream);

		auto t_end = std::chrono::high_resolution_clock::now();
		float total_inf = std::chrono::duration<float, std::milli>(t_end - t_start).count();
		std::cout << "Inference take: " << total_inf << " ms." << std::endl;

		auto p_start = std::chrono::high_resolution_clock::now();

		auto vv_boxes = PostProcess(vec_Mat, p_imgs, m_v_out_buf[0].data(), m_out_size[1]);

		auto p_end = std::chrono::high_resolution_clock::now();
		float total_p = std::chrono::duration<float, std::milli>(p_end - p_start).count();
		//std::cout << "PostProcess take: " << total_p << " ms." << std::endl;

		vv_results = vv_boxes;
	}
	
	std::vector<std::vector<DetectRes>> CKeyPoint::PostProcess(const std::vector<cv::Mat>& vec_Mat, st_dlm_data* p_imgs, float* output, const int& outSize)
	{
		int batch = m_info.batch_size;
		int model_width = m_info.image_width;
		int model_height = m_info.image_height;
		int model_channel = m_info.input_channel;

		std::vector<std::vector<DetectRes>> vec_result;
		int index = 0;
		for (const cv::Mat& src_img : vec_Mat)
		{
			std::vector<DetectRes> result;
			float ratio = float(src_img.cols) / float(model_width) > float(src_img.rows) / float(model_height) ? 
				float(src_img.cols) / float(model_width) :
				float(src_img.rows) / float(model_height);

			float* out = output + index * outSize;

			for (int i = 0; i < m_num_anchors; i++) {

				float scale = __min(float(model_width) / float(src_img.cols), float(model_height) / float(src_img.rows));
				int padd_h = (model_height - src_img.rows * scale) / 2;
				int padd_w = (model_width - src_img.cols * scale) / 2;

				DetectRes box;
				box.classes = 0;
				box.prob = out[i + m_num_anchors * 4];
				if (box.prob < p_imgs[index].prob_threshold)
					continue;
				box.x = (out[i] - padd_w) * ratio;
				box.y = (out[i + m_num_anchors] - padd_h) * ratio;
				box.w = out[i + m_num_anchors * 2] * ratio;
				box.h = out[i + m_num_anchors * 3] * ratio;

				// TODO: multi keypoint
				box.kps_s = out[i + m_num_anchors * 7];
				if (box.kps_s < p_imgs[index].kpt_threshold)
					continue;
				box.kps_x = (out[i + m_num_anchors * 5] - padd_w) * ratio;
				box.kps_y = (out[i + m_num_anchors * 6] - padd_h) * ratio;

				result.push_back(box);
			}

			NmsDetect(result,p_imgs[index].nms_threshold);
			vec_result.push_back(result);
			index++;
		}
		return vec_result;
	}

	void CKeyPoint::NmsDetect(std::vector<DetectRes>& detections, float nms_threshold)
	{
		sort(detections.begin(), detections.end(), [=](const DetectRes &left, const DetectRes &right) {
			return left.prob > right.prob;
		});

		for (int i = 0; i < (int)detections.size(); i++)
			for (int j = i + 1; j < (int)detections.size(); j++)
			{
				if (detections[i].classes == detections[j].classes)
				{
					float iou = IOUCalculate(detections[i], detections[j]);
					if (iou > nms_threshold)
						detections[j].prob = 0;
				}
			}

		detections.erase(std::remove_if(detections.begin(), detections.end(), [](const DetectRes &det)
		{ return det.prob == 0; }), detections.end());
	}

	float CKeyPoint::IOUCalculate(const DetectRes& det_a, const DetectRes& det_b)
	{
		cv::Point2f center_a(det_a.x, det_a.y);
		cv::Point2f center_b(det_b.x, det_b.y);
		cv::Point2f left_up(__min(det_a.x - det_a.w / 2, det_b.x - det_b.w / 2),
			__min(det_a.y - det_a.h / 2, det_b.y - det_b.h / 2));
		cv::Point2f right_down(__max(det_a.x + det_a.w / 2, det_b.x + det_b.w / 2),
			__max(det_a.y + det_a.h / 2, det_b.y + det_b.h / 2));
		float distance_d = (center_a - center_b).x * (center_a - center_b).x + (center_a - center_b).y * (center_a - center_b).y;
		float distance_c = (left_up - right_down).x * (left_up - right_down).x + (left_up - right_down).y * (left_up - right_down).y;
		float inter_l = det_a.x - det_a.w / 2 > det_b.x - det_b.w / 2 ? det_a.x - det_a.w / 2 : det_b.x - det_b.w / 2;
		float inter_t = det_a.y - det_a.h / 2 > det_b.y - det_b.h / 2 ? det_a.y - det_a.h / 2 : det_b.y - det_b.h / 2;
		float inter_r = det_a.x + det_a.w / 2 < det_b.x + det_b.w / 2 ? det_a.x + det_a.w / 2 : det_b.x + det_b.w / 2;
		float inter_b = det_a.y + det_a.h / 2 < det_b.y + det_b.h / 2 ? det_a.y + det_a.h / 2 : det_b.y + det_b.h / 2;
		if (inter_b < inter_t || inter_r < inter_l)
			return 0;
		float inter_area = (inter_b - inter_t) * (inter_r - inter_l);
		float union_area = det_a.w * det_a.h + det_b.w * det_b.h - inter_area;
		if (union_area == 0)
			return 0;
		else
			return inter_area / union_area - distance_d / distance_c;
	}
}